/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "aicpu_cust_kernel_builder.h"

#include <fstream>
#include "cce/taskdown_common.hpp"
#include "aicpu/common/aicpu_task_struct.h"

#include "util/constant.h"
#include "error_code/error_code.h"
#include "util/log.h"
#include "util/util.h"
#include "config/ops_json_file.h"

namespace aicpu {
KernelBuilderPtr AicpuCustKernelBuilder::instance_ = nullptr;

KernelBuilderPtr AicpuCustKernelBuilder::Instance()
{
    static std::once_flag flag;
    std::call_once(flag, [&]() {
        instance_.reset(new (std::nothrow) AicpuCustKernelBuilder);
    });
    return instance_;
}

ge::Status AicpuCustKernelBuilder::CalcOpRunningParam(const ge::Node &node) const
{
    const std::string node_name = node.GetName();
    const std::string node_type = node.GetType();
    AICPUE_LOGI("AICPUKernel's op %s[%s] run CalcOpRunningParam",
        node_name.c_str(), node_type.c_str());

    std::shared_ptr<ge::OpDesc> op_desc_ptr = node.GetOpDesc();
    AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, ErrorCode::INPUT_PARAM_NULL)

    // workspace size
    // workspace size unit kb. The maximum size of workspace_size is 1G, no need to check overflow
    int64_t workspace_size = 0;
    (void)ge::AttrUtils::GetInt(op_desc_ptr, "workspaceSize", workspace_size);
    op_desc_ptr->SetWorkspaceBytes({workspace_size});
    AICPU_CHECK_RES_WITH_LOG(KernelBuilder::SetOutPutsSize(op_desc_ptr),
        "Call SetOutPutsSize function failed, op[%s]", node_name.c_str())

    AICPU_CHECK_FALSE_EXEC(ge::AttrUtils::SetListBool(op_desc_ptr, kWorkspaceReuseFlag, {true}),
        AICPU_REPORT_CALL_ERROR("Set workspace memory reuse flag failed, op[%s]",
            node_name.c_str());
        return ErrorCode::ADD_ATTR_FAILED)

    AICPUE_LOGI("AICPUKernel's op %s[%s] run CalcOpRunningParam success. workspace total size is %lld",
                node_name.c_str(), node_type.c_str(), workspace_size);
    return ge::SUCCESS;
}

ge::Status AicpuCustKernelBuilder::GenerateTask(const ge::Node &node, const ge::RunContext &run_context,
                                                std::vector<domi::TaskDef> &tasks)
{
    std::string op_name = node.GetName();
    std::string op_type = node.GetType();
    AICPUE_LOGI("AICPUKernel's op %s[%s] run GenerateTask",
        op_name.c_str(), op_type.c_str());
    auto op_desc_ptr = node.GetOpDesc();
    AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)
    AICPU_CHECK_NOTNULL_ERRCODE(run_context.model, INPUT_PARAM_NULL)
    AICPU_CHECK_NOTNULL_ERRCODE(run_context.stream, INPUT_PARAM_NULL)

    domi::TaskDef task_def;
    task_def.set_type(RT_MODEL_TASK_KERNEL);
    // no need to set streamID for task_def, ge will reallocatee stream
    domi::KernelDef *kernel_def = task_def.mutable_kernel();
    AICPU_CHECK_NOTNULL(kernel_def)

    // create prame_base and launch kernel
    AICPU_CHECK_RES_WITH_LOG(
        BuildAndLaunchKernel(op_desc_ptr, run_context, kernel_def),
        "Call BuildAndLaunchKernel function failed, op[%s]", op_name.c_str())

    (void)tasks.emplace_back(task_def);
    AICPUE_LOGI("AICPUKernel's op %s[%s] run GenerateTask success",
        op_name.c_str(), op_type.c_str());
    return ge::SUCCESS;
}

ge::Status AicpuCustKernelBuilder::BuildAndLaunchKernel(const ge::OpDescPtr &op_desc_ptr,
                                                        const ge::RunContext &run_context,
                                                        domi::KernelDef *&kernel_def) const
{
    std::vector<char> task_ext_info;
    AICPU_CHECK_RES_WITH_LOG(MakeBaseExtInfo(op_desc_ptr, task_ext_info),
        "Call MakeTaskExtInfo function failed, op[%s]",
        op_desc_ptr->GetName().c_str());
    if (task_ext_info.size() == 0) {
        kernel_def->set_kernel_ext_info_size(0);
    } else {
        kernel_def->set_kernel_ext_info(reinterpret_cast<void *>(task_ext_info.data()), task_ext_info.size());
        kernel_def->set_kernel_ext_info_size(task_ext_info.size());
    }

    // set kernel so name's len and offset addr
    std::string kernel_so_name;
    (void)ge::AttrUtils::GetStr(op_desc_ptr, "kernelSo", kernel_so_name);
    if (kernel_so_name.empty()) {
        AICPU_REPORT_INNER_ERROR("attr kernelSo is empty, op[%s].",
            op_desc_ptr->GetName().c_str());
        return STR_IS_EMPTY;
    }

    // set kernel function name's len and offset addr
    std::string func_name;
    (void)ge::AttrUtils::GetStr(op_desc_ptr, "funcName", func_name);
    if (func_name.empty()) {
        AICPU_REPORT_INNER_ERROR("funcName of empty, op[%s].",
            op_desc_ptr->GetName().c_str())
        return STR_IS_EMPTY;
    }

    // calculate param_len: AicpuParamHead.len + io_addrs_size + notify_id.len + customizedAttr.len
    bool flag_async = false;
    (void)ge::AttrUtils::GetBool(op_desc_ptr, "paramLen", flag_async);
    uint32_t param_len = flag_async ? static_cast<uint32_t>(sizeof(AicpuParamHead) + sizeof(uint32_t)) :
                                           static_cast<uint32_t>(sizeof(AicpuParamHead));
    // get input and output total number, no need to check overflow
    uint32_t io_addrs_num = static_cast<uint32_t>(op_desc_ptr->GetInputsSize() + op_desc_ptr->GetOutputsSize());
    // get input and output addrs size, no need to check overflow
    uint32_t io_addrs_size = io_addrs_num * static_cast<uint32_t>(sizeof(uint64_t));
    // refresh param_len, no need to check overflow
    param_len += io_addrs_size;

    ge::GeAttrValue::BYTES bytes;
    bool has_customized_attr = ge::AttrUtils::GetZeroCopyBytes(op_desc_ptr, kCustomizedOpDef, bytes);
    // When it's aicpu customized ops, get customized attr
    if (has_customized_attr) {
        CHECK_UINT32_ADD_OVERFLOW(param_len, sizeof(uint32_t), ErrorCode::DATA_OVERFLOW,
            "Overflow when param total bytes[%u] add 4bytes, op[%s]",
            param_len, op_desc_ptr->GetName().c_str())
        param_len += sizeof(uint32_t);
        // Customized attr length must be less than UINT32_MAX, no need to check overflow
        uint32_t customized_attr_len = static_cast<uint32_t>(bytes.GetSize());
        CHECK_UINT32_ADD_OVERFLOW(param_len, customized_attr_len, ErrorCode::DATA_OVERFLOW,
            "Overflow when calculate total bytes of task param[%u] and custom "
            "attr[%u], op[%s]", param_len, customized_attr_len,
            op_desc_ptr->GetName().c_str())
        param_len += customized_attr_len;
    }

    uint32_t ext_info_length = task_ext_info.size();
    uint64_t ext_info_addrs = 0;
    std::string task_args;
    // Create task_args: AicpuParamHead + ioAddrs + notify_id + customizedAttr
    AicpuParamHead param_head = {param_len, io_addrs_num, ext_info_length, ext_info_addrs};
    task_args.append(reinterpret_cast<const char *>(&param_head), sizeof(AicpuParamHead));
    // TaskArgs append ioAddrs
    task_args.append(io_addrs_size, ' ');
    // TaskArgs append notify_id for Async op
    if (flag_async) {
        uint32_t notify_id = 0;
        task_args.append(reinterpret_cast<const char *>(&notify_id), sizeof(uint32_t));
    }

    // When it's aicpu customized ops, task_args should append customized attr
    if (has_customized_attr) {
        const uint8_t *node_def_data = bytes.GetData();
        AICPU_CHECK_NOTNULL(node_def_data)
        size_t customized_attr_len = bytes.GetSize();
        uint32_t node_def_len = static_cast<uint32_t>(customized_attr_len);
        task_args.append(reinterpret_cast<const char *>(&node_def_len), sizeof(uint32_t));
        task_args.append(reinterpret_cast<const char *>(node_def_data), customized_attr_len);
    }

    kernel_def->set_args(task_args.data(), param_len);
    kernel_def->set_args_size(param_len);
    kernel_def->set_so_name(kernel_so_name);
    kernel_def->set_kernel_name(func_name);

    domi::KernelContext *context = kernel_def->mutable_context();
    AICPU_CHECK_NOTNULL(context)
    context->set_kernel_type(cce::ccKernelType::CUST_AI_CPU);
    context->set_op_index(op_desc_ptr->GetId());

    AICPUE_LOGI("task_args length is %zu, param_len is %u", task_args.length(), param_len);
    return ge::SUCCESS;
}

FACTORY_KERNEL_BUILDER_CLASS_KEY(AicpuCustKernelBuilder, "CUSTAICPUBuilder")
}
